import random

import numpy as np
import torchvision
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import transforms

from _dynamic_dataset.TrainDataset import TrainSubset
from rootPath import project_path
from server_client.Client import Client


def splitTrainDataByNums(orderList, trainDataset, seed=0):
    # 设置随机种子
    random.seed(seed)
    # 存储数据索引
    datasetList = []
    # 初始化一个字典来存储每个类别的数据索引
    data_indices = {cls: [] for cls in range(10)}
    # 记录每个类别的数据索引
    for i, (image, label) in enumerate(trainDataset):
        data_indices[label].append(i)

    indexNum = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    for order in orderList:
        indices = []
        for cls in range(10):
            cls_indices = data_indices[cls]
            if indexNum[cls] + order[cls] > len(cls_indices):
                raise IndexError("数组下标越界，请检查客户端各类数据总数")
            sampled_indices = cls_indices[indexNum[cls]:indexNum[cls] + order[cls]]
            indexNum[cls] += order[cls]
            indices.extend(sampled_indices)
        # 进行一个indice随机化，但不会改变每个子集所含的样本
        random.shuffle(indices)
        datasetList.append(indices)

    return datasetList


def addNoisyToLabelInTrainDataset(train_dataset, indices, ratio):
    """

    :param train_dataset: cifar-10数据集
    :param indices: 可能被修改的数据的序号
    :param ratio: 比例
    :return:
    """
    indices_to_change = random.sample(indices, int(ratio * len(indices)))
    # 更改指定样本的标签
    for index in indices_to_change:
        # 获取原始标签
        original_label = train_dataset.targets[index]
        # 生成新的随机标签
        new_label = random.randint(0, 9)
        while new_label == original_label:
            new_label = random.randint(0, 9)
        # 替换标签
        train_dataset.targets[index] = new_label
    return indices_to_change


def addNoisyToFeaturesInTrainDataset(train_dataset, indices, ratio, noise_level=0.1):
    """
    :param train_dataset: cifar-10数据集
    :param indices: 可能被修改的数据的序号
    :param ratio: 比例
    :param noise_level: 噪音强度，值越大噪音越明显
    :return: 修改的索引列表
    """
    indices_to_change = random.sample(indices, int(ratio * len(indices)))

    for index in indices_to_change:
        # 获取原始图像
        original_image = train_dataset.data[index]
        # 生成高斯噪音
        noise = np.random.normal(0, noise_level * 255, original_image.shape).astype(np.uint8)
        # 添加噪音
        noisy_image = np.clip(original_image + noise, 0, 255)
        # 替换图像
        train_dataset.data[index] = noisy_image

    return indices_to_change


# 显示添加噪声后的图像
def show_images(images, labels):
    fig, axs = plt.subplots(2, 5, figsize=(15, 6))
    axs = axs.flatten()
    for i, ax in enumerate(axs):
        if i < len(images):
            ax.imshow(np.transpose(images[i], (0, 1, 2)))  # 将通道维度移到最后
            ax.set_title(f"Label: {labels[i]}")
            ax.axis('off')
    plt.tight_layout()
    plt.show()


def get_SDSS_clients(clientConf, model, device, train_transform):
    clients = []
    numsList = [
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
    ]
    clientNum = len(numsList)
    # 测试集
    train_dataset = torchvision.datasets.CIFAR10(root=project_path + '/data/', train=True, download=True,
                                                 transform=train_transform)
    datasetIndices = splitTrainDataByNums(orderList=numsList, trainDataset=train_dataset)
    for i in range(clientNum):
        trainSubDataset = TrainSubset(train_dataset, datasetIndices[i])
        clients.append(
            Client(clientConf, model, device,
                   DataLoader(trainSubDataset, shuffle=True, batch_size=clientConf['batch_size'], num_workers=2), i))
        print(f"客户端{i}创建成功")
    print("客户端全部初始化成功！")
    return clients

def get_DDSS_clients(clientConf, model, device, train_transform):
    clients = []
    numsList = [
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 200, 9: 200},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 200, 9: 200},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 200, 9: 200},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 200, 9: 200},
        {0: 800, 1: 800, 2: 800, 3: 800, 4: 800, 5: 800, 6: 800, 7: 800, 8: 1000, 9: 1000},
    ]
    clientNum = len(numsList)
    # 测试集
    train_dataset = torchvision.datasets.CIFAR10(root=project_path + '/data/', train=True, download=True,
                                                 transform=train_transform)
    datasetIndices = splitTrainDataByNums(orderList=numsList, trainDataset=train_dataset)
    for i in range(clientNum):
        trainSubDataset = TrainSubset(train_dataset, datasetIndices[i])
        clients.append(
            Client(clientConf, model, device,
                   DataLoader(trainSubDataset, shuffle=True, batch_size=clientConf['batch_size'], num_workers=2), i))
        print(f"客户端{i}创建成功")
    print("客户端全部初始化成功！")
    return clients

def get_SDDS_clients(clientConf, model, device, train_transform):
    clients = []
    numsList = [
        {0: 500, 1: 500, 2: 500, 3: 500, 4: 500, 5: 500, 6: 500, 7: 500, 8: 500, 9: 500},
        {0: 500, 1: 500, 2: 500, 3: 500, 4: 500, 5: 500, 6: 500, 7: 500, 8: 500, 9: 500},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 2000, 1: 2000, 2: 2000, 3: 2000, 4: 2000, 5: 2000, 6: 2000, 7: 2000, 8: 2000, 9: 2000},
    ]
    clientNum = len(numsList)
    # 测试集
    train_dataset = torchvision.datasets.CIFAR10(root=project_path + '/data/', train=True, download=True,
                                                 transform=train_transform)
    datasetIndices = splitTrainDataByNums(orderList=numsList, trainDataset=train_dataset)
    for i in range(clientNum):
        trainSubDataset = TrainSubset(train_dataset, datasetIndices[i])
        clients.append(
            Client(clientConf, model, device,
                   DataLoader(trainSubDataset, shuffle=True, batch_size=clientConf['batch_size'], num_workers=2), i))
        print(f"客户端{i}创建成功")
    print("客户端全部初始化成功！")
    return clients


def get_NFSS_clients(clientConf, model, device, train_transform, noisy_list):
    clients = []
    numsList = [
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
    ]
    clientNum = len(numsList)
    # 训练集
    train_dataset = torchvision.datasets.CIFAR10(root=project_path + '/data/', train=True, download=True,
                                                 transform=train_transform)
    datasetIndices = splitTrainDataByNums(orderList=numsList, trainDataset=train_dataset)
    for i in range(clientNum):
        addNoisyToFeaturesInTrainDataset(train_dataset, datasetIndices[i], noisy_list[i])
        trainSubDataset = TrainSubset(train_dataset, datasetIndices[i])
        clients.append(
            Client(clientConf, model, device,
                   DataLoader(trainSubDataset, shuffle=True, batch_size=clientConf['batch_size'], num_workers=2), i))
        print(f"客户端{i}创建成功")
    print("客户端全部初始化成功！")
    return clients


def get_NLSS_clients(clientConf, model, device, train_transform, noisy_list):
    clients = []
    numsList = [
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
        {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
    ]
    clientNum = len(numsList)
    # 训练集
    train_dataset = torchvision.datasets.CIFAR10(root=project_path + '/data/', train=True, download=True,
                                                 transform=train_transform)
    datasetIndices = splitTrainDataByNums(orderList=numsList, trainDataset=train_dataset)
    for i in range(clientNum):
        addNoisyToLabelInTrainDataset(train_dataset, datasetIndices[i], noisy_list[i])
        trainSubDataset = TrainSubset(train_dataset, datasetIndices[i])
        clients.append(
            Client(clientConf, model, device,
                   DataLoader(trainSubDataset, shuffle=True, batch_size=clientConf['batch_size'], num_workers=2), i))
        print(f"客户端{i}创建成功")
    print("客户端全部初始化成功！")
    return clients

# 示例使用
if __name__ == '__main__':
    # 加载CIFAR-10数据集
    transform = transforms.Compose([transforms.ToTensor()])
    train_dataset = datasets.CIFAR10(root=project_path + '/data/', train=True, download=True, transform=transform)

    # 添加噪声并展示结果
    ratio = 1  # 10%的样本添加噪声
    indices = list(range(10))

    # 添加噪声到特征
    noisy_indices = addNoisyToFeaturesInTrainDataset(train_dataset, indices, ratio, noise_level=0.05)
    noisy_indices = sorted(noisy_indices)
    # 获取添加噪声的图像和标签
    noisy_images = train_dataset.data[noisy_indices]
    noisy_labels = [train_dataset.classes[train_dataset.targets[i]] for i in noisy_indices]

    # 显示添加噪声后的图像
    show_images(noisy_images, noisy_labels)
